package com.fudan.run.trans;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.UUID;

import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.function.Function2;
import org.apache.spark.broadcast.Broadcast;
import org.apache.spark.util.LongAccumulator;

import com.fasterxml.jackson.core.type.TypeReference;
import com.fudan.cfg.trans.RpcTransDef;
import com.fudan.run.JobRunner;
import com.fudan.run.VariableRddMap;
import com.fudan.run.ctx.annotation.TransAdapter;
import com.lobstar.xrpc.client.main.MainClient;
import com.lobstar.xrpc.tools.Fields;

@TransAdapter("rpc")
public class RpcTrans extends TransHandler<RpcTransDef>{

	@Override
	public void trans(JobRunner jobRunner,RpcTransDef transNode, VariableRddMap variableMap) {
		String target = transNode.getTarget();
		JavaRDD<Map<String, Object>> targetRdd = variableMap.getRdd(target);
		List<String> servers = transNode.getServers();
		Broadcast<Integer> serverSize = jobRunner.getJavaSparkContext().broadcast(servers.size());

		JavaRDD<Map<String, Object>> after = targetRdd.mapPartitionsWithIndex(new Function2<Integer, Iterator<Map<String,Object>>, Iterator<Map<String,Object>>>() {
			@Override
			public Iterator<Map<String, Object>> call(Integer i, Iterator<Map<String, Object>> l) throws Exception {
				Integer size = serverSize.getValue();
				int index = (i+1)%size;
				String server = servers.get(index);
				String[] ipPort = server.split("[:]");
				List<Map<String,Object>> list = new ArrayList<>();
				MainClient client = MainClient.build(UUID.randomUUID().toString(), ipPort[0], Integer.parseInt(ipPort[1]), null).blockStart();
				while(l.hasNext()) {
					Map<String, Object> map = l.next();
					Map<String,Object> r = client.asyncCall(transNode.getMethod(), Fields.build().field("map", map), Map.class);
					list.add(r);
				}
				return list.iterator();
			}
		},false);
//		JavaRDD<Map<String, Object>> after = targetRdd.mapPartitions(l->{
//			//String server = servers.get(index);
//			//String[] ipPort = server.split("[:]");
//			List<Map<String,Object>> list = new ArrayList<>();
//			MainClient client = MainClient.build(UUID.randomUUID().toString(), ipPort[0], Integer.parseInt(ipPort[1]), null).blockStart();
//			while(l.hasNext()) {
//				Map<String, Object> map = l.next();
//				Map<String,Object> r = client.asyncCall(transNode.getMethod(), Fields.build().field("map", map), Map.class);
//				list.add(r);
//			}
//			return list.iterator();
//		});
		variableMap.addRdd(transNode.getName(), after);
	}

}
